# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
import random
from typing import Union, List
from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig
 
 
 
@GENERATOR_REGISTRY.register("ascend_generate_aclnn_cross_entropy_loss_grad")
class CtcLossBackwardGenerator(CaseGenerator):
 
    def __init__(self, config):
        super().__init__(config)
        self.dtype = None
        self.logProb_shape = []
 
    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]]
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:
        '''
        当参数之间有相互依赖关系时，需要覆写此函数来约束生成的输入参数信息
        :param index: 用例参数列表的下标，0：表示第0个参数
        :param input_case: 随机生成的用例信息
        :return: 修改后需要返回的输入信息对象
        '''
 
        if index == 0:
            self.dtype = input_case.dtype
 
        if index == 1:
            input_case.dtype = self.dtype
 
        if index == 4:
            input_case.dtype = self.dtype
 
        if index == 5:
            input_case.dtype = self.dtype
 
        return input_case
 
    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
 
        reduction = case_config.inputs[6]
        if 0 in case_config.inputs[1].shape: ## 单独约束空tensor:按照约束条件空tensor只能N==0
            case_config.inputs[1].shape[0] = 0
            case_config.inputs[1].shape[1] = 131073
            [N, C] = case_config.inputs[1].shape
            if reduction.range_values in ["mean", "sum"]:
                case_config.inputs[0].shape = []
            else:
                case_config.inputs[0].shape[0] = N
            case_config.inputs[2].shape[0] = N
            case_config.inputs[2].range_values = [0, C - 1]
            case_config.inputs[3].shape[0] = C
            case_config.inputs[7].range_values = random.randint(-100, C - 1)
        elif case_config.inputs[1].shape == []: ## 单独约束标量tensor
            pass
        else:
            [N, C] = case_config.inputs[1].shape
            if N >= 2147483649: # 上边界用例
                N = 190000 #算子约束logProb第零维N需满足N<200000
            if reduction.range_values in ["mean", "sum"]:
                case_config.inputs[0].shape = []
            else:
                case_config.inputs[0].shape[0] = N
            case_config.inputs[2].shape[0] = N
            case_config.inputs[2].range_values = [0, C - 1]
            case_config.inputs[3].shape[0] = C
            case_config.inputs[7].range_values = random.randint(-100, C - 1)
        return case_config